#include <iostream>
#include "Ws2tcpip.h"
#include "Winsock2.h"

using std::cin;
using std::cout;
using std::cerr;
using std::endl;

#pragma comment(lib, "ws2_32.lib")

int SOCKADDR_IN_SIZE = sizeof(SOCKADDR_IN);
const u_short DEFAULT_PORT = 4000;
const size_t MSG_BUF_SIZE = 512;
const size_t IP_BUF_SIZE = 256;

int main() {
    WSADATA wsa_data;
    SOCKET sock_serv = INVALID_SOCKET;
    SOCKADDR_IN addr_serv, addr_clt;
    char ip_buf[IP_BUF_SIZE];
    char msg_buf[MSG_BUF_SIZE];
    int ret_val = 0;
    //
    ret_val = WSAStartup(MAKEWORD(2, 2), &wsa_data);
    if (ret_val != 0) {
        cerr << "WSAStartup() function failed with error: " << WSAGetLastError() << "\n";
        system("pause");
        return 1;
    }
    //
    //SecureZeroMemory(&addr_serv, SOCKADDR_IN_SIZE);
    addr_serv.sin_family = AF_INET;
    addr_serv.sin_port = htons(DEFAULT_PORT);
    addr_serv.sin_addr.S_un.S_addr = ADDR_ANY;
    //
    sock_serv = socket(addr_serv.sin_family, SOCK_DGRAM, IPPROTO_UDP);
    if (sock_serv == INVALID_SOCKET) {
        cerr << "socket() function failed with error: " << WSAGetLastError() << "\n";
        WSACleanup();
        
        system("pause");
        return 1;
    }
    //
    ret_val = bind(sock_serv, (SOCKADDR*)&addr_serv, SOCKADDR_IN_SIZE);
    if (ret_val != 0) {
        cerr << "bind() function failed with error: " << WSAGetLastError() << "\n";
        system("pause");
        return 1;
    }
    cout << "A UDP server has started successfully..." << endl;
    //
    while (true) {
        SecureZeroMemory(msg_buf, MSG_BUF_SIZE);
        ret_val = recvfrom(sock_serv, msg_buf, MSG_BUF_SIZE, 0, (SOCKADDR*)&addr_clt, &SOCKADDR_IN_SIZE);
        if (ret_val > 0) {
            inet_ntop(addr_clt.sin_family, &addr_clt, ip_buf, IP_BUF_SIZE);
            cout << "message from client " << ip_buf << ": " << msg_buf << endl;
            cout << "enter the message to response: ";
            cin.getline(msg_buf, MSG_BUF_SIZE);
            sendto(sock_serv, msg_buf, strlen(msg_buf), 0, (SOCKADDR*)&addr_clt, SOCKADDR_IN_SIZE);
        }
        else if (ret_val == 0) {
            cout << "connection is closed..." << endl;
        }
        else {
            cerr << "recvfrom() function failed with error: " << WSAGetLastError() << "\n";
            closesocket(sock_serv);
            WSACleanup();
            system("pause");
            return 1;
        }
    }
    ret_val = shutdown(sock_serv, SD_BOTH);
    if (ret_val == SOCKET_ERROR) {
        cerr << "shutdown() function failed with error: " << WSAGetLastError() << "\n";
        closesocket(sock_serv);
        WSACleanup();
        system("pause");
        return 1;
    }
    closesocket(sock_serv);
    WSACleanup();
    cout << "server shutdown..." << endl;
    system("pause");
    return 0;
}